import streamlit as st
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import numpy as np
import os


# 定义AlexNet模型类（与训练时相同）
class AlexNet(nn.Module):
    def __init__(self, num_classes=8, dropout_rate=0.5):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
            nn.MaxPool2d(kernel_size=3, stride=2),

            nn.Conv2d(96, 256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
            nn.MaxPool2d(kernel_size=3, stride=2),

            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(p=dropout_rate),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=dropout_rate),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


# 加载模型
@st.cache_resource
def load_model():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 检查模型文件是否存在
    if not os.path.exists('alexnet_model.pth'):
        st.error("模型文件 'alexnet_model.pth' 未找到。请先运行 train.py 训练模型。")
        return None, None

    # 加载模型
    checkpoint = torch.load('alexnet_model.pth', map_location=device)
    model = AlexNet(num_classes=len(checkpoint['class_to_idx']))
    model.load_state_dict(checkpoint['model_state_dict'])
    model.to(device)
    model.eval()

    return model, checkpoint


# 图像预处理
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((227, 227)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)


# 预测函数
def predict(image, model, device, idx_to_class):
    image_tensor = preprocess_image(image).to(device)

    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)
        confidence, predicted = torch.max(probabilities, 1)

    predicted_class = idx_to_class[predicted.item()]
    confidence = confidence.item()

    return predicted_class, confidence


# 主应用
def main():
    st.set_page_config(
        page_title="AlexNet图像分类器",
        page_icon="🖼️",
        layout="wide"
    )

    st.title("🖼️ AlexNet图像分类器")
    st.markdown("使用在Natural Images数据集上训练的AlexNet模型进行图像分类")

    # 加载模型
    with st.spinner("正在加载模型..."):
        model, checkpoint = load_model()

    if model is None:
        st.stop()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    idx_to_class = checkpoint['idx_to_class']

    st.sidebar.title("设置")

    # 图像上传
    uploaded_file = st.sidebar.file_uploader(
        "上传图像",
        type=['png', 'jpg', 'jpeg'],
        help="支持PNG、JPG、JPEG格式"
    )


    col1, col2 = st.columns(2)

    with col1:
        st.subheader("输入图像")

        if uploaded_file is not None:
            image = Image.open(uploaded_file).convert('RGB')
            st.image(image, caption="上传的图像", use_column_width=True)
        else:
            st.info("请上传图像或选择使用示例图像")
            image = None

    with col2:
        st.subheader("分类结果")

        if image is not None:
            with st.spinner("正在分类..."):
                predicted_class, confidence = predict(image, model, device, idx_to_class)

            st.success(f"预测类别: **{predicted_class}**")
            st.info(f"置信度: **{confidence:.2%}**")

            # 显示所有类别的概率
            st.subheader("所有类别概率")

            image_tensor = preprocess_image(image).to(device)
            with torch.no_grad():
                outputs = model(image_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]

            # 创建概率条形图
            classes = list(idx_to_class.values())
            probs = probabilities.cpu().numpy()

            # 按概率排序
            sorted_indices = np.argsort(probs)[::-1]
            sorted_classes = [classes[i] for i in sorted_indices]
            sorted_probs = [probs[i] for i in sorted_indices]

            # 显示概率条
            for i, (cls, prob) in enumerate(zip(sorted_classes, sorted_probs)):
                st.write(f"{cls}: {prob:.2%}")
                st.progress(float(prob))

    # 模型信息
    st.sidebar.markdown("---")
    st.sidebar.subheader("模型信息")
    st.sidebar.write(f"设备: {device}")
    st.sidebar.write(f"类别数量: {len(idx_to_class)}")
    st.sidebar.write("支持的类别:")
    for idx, class_name in idx_to_class.items():
        st.sidebar.write(f"- {class_name}")


if __name__ == "__main__":
    main()